package rsa

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"gitee.com/zackeus/go-boot/tools/errorx"
	"github.com/go-jose/go-jose/v4"
)

// GenerateRSA 生成 RSA 密钥对
func GenerateRSA(alg jose.SignatureAlgorithm) (*rsa.PrivateKey, *rsa.PublicKey, error) {
	/* RS256/RS384 推荐 2048 位密钥 */
	keyBits := 2048
	switch alg {
	case jose.RS512, jose.PS512:
		keyBits = 4096
	}
	privKey, err := rsa.GenerateKey(rand.Reader, keyBits)
	if err != nil {
		return nil, nil, errorx.Wrap(err, "生成 RSA 密钥失败")
	}
	return privKey, &privKey.PublicKey, nil
}

// GenerateRSAPem 生成 RSA 密钥对
func GenerateRSAPem(alg jose.SignatureAlgorithm) ([]byte, []byte, error) {
	privKey, pubKey, err := GenerateRSA(alg)
	if err != nil {
		return nil, nil, err
	}

	/* 编码私钥 (PKCS#1 格式) */
	privPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	/* 编码公钥 (PKIX 格式) */
	pubBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, nil, errorx.Wrap(err, "编码 RSA 公钥失败")
	}
	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubBytes,
	})

	return privPEM, pubPEM, nil
}

// LoadRSAPrivateKey 将 PEM 格式的 RSA 私钥字符串解析为 *rsa.PrivateKey
// privPemBytes 是通过 GenerateRSAPem 方法生成的私钥字节流
func LoadRSAPrivateKey(privPemBytes []byte) (*rsa.PrivateKey, error) {
	// 1. 解码 PEM 数据
	// pem.Decode 会从输入的字节流中提取第一个 PEM 块
	block, _ := pem.Decode(privPemBytes)
	if block == nil {
		return nil, errors.New("未能解码 PEM 数据，输入可能不是有效的 PEM 格式")
	}

	// 2. 检查 PEM 块的类型是否为 RSA PRIVATE KEY
	// 这与我们在 GenerateRSAPem 中设置的 Type "RSA PRIVATE KEY" 相对应
	if block.Type != "RSA PRIVATE KEY" {
		return nil, errors.New(fmt.Sprintf("不支持的 PEM 块类型: %s，期望: RSA PRIVATE KEY", block.Type))
	}

	// 3. 解析 DER 编码的私钥
	// x509.ParsePKCS1PrivateKey 用于解析 PKCS#1 格式的私钥
	privKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, errorx.Wrap(err, "解析 RSA 私钥失败")
	}

	// 4. 验证私钥 可以确保加载的私钥是有效的
	err = privKey.Validate()
	if err != nil {
		return nil, errorx.Wrap(err, "加载的 RSA 私钥无效")
	}
	return privKey, nil
}

// LoadRSAPublicKey 将 PEM 格式的 RSA 公钥字符串解析为 *rsa.PublicKey
// pubPemBytes 是通过 GenerateRSAPem 方法生成的公钥字节流
func LoadRSAPublicKey(pubPemBytes []byte) (*rsa.PublicKey, error) {
	// 1. 解码 PEM 数据
	block, _ := pem.Decode(pubPemBytes)
	if block == nil {
		return nil, errors.New("未能解码 PEM 数据，输入可能不是有效的 PEM 格式")
	}

	// 2. 检查 PEM 块的类型是否为 RSA PUBLIC KEY
	// 这与我们在 GenerateRSAPem 中设置的 Type "RSA PUBLIC KEY" 相对应
	if block.Type != "RSA PUBLIC KEY" {
		return nil, errors.New(fmt.Sprintf("不支持的 PEM 块类型: %s，期望: RSA PUBLIC KEY", block.Type))
	}

	// 3. 解析 DER 编码的公钥
	// x509.ParsePKIXPublicKey 可以解析多种格式的公钥，返回一个 interface{}
	pubKeyInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, errorx.Wrap(err, "解析 RSA 公钥失败")
	}

	// 4. 将解析出的公钥断言为 *rsa.PublicKey 类型
	// 因为我们明确知道这是一个 RSA 公钥，所以可以安全地进行类型断言
	pubKey, ok := pubKeyInterface.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New(fmt.Sprintf("解析出的公钥类型不是 RSA，而是: %T", pubKeyInterface))
	}

	return pubKey, nil
}
